package edu.cmu.cs211.compression.huffman;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;

import java.util.Comparator;
import java.util.Map;
import java.util.Stack;
import java.util.Vector;

import edu.cmu.cs211.compression.io.BitReader;
import edu.cmu.cs211.compression.io.BitWriter;
import edu.cmu.cs211.compression.io.InputStreamBitReader;
import edu.cmu.cs211.compression.util.MyPriorityQueue;

/**
 * Represents the Huffman code. The class supports building the Huffman tree
 * based either on the frequency of given symbols or by reading the tree from an
 * on-disk format. It also supports emitting the code word for a given symbol or
 * reading a symbol based on the code word. Lastly, it is able to write an
 * on-disk representation of the Huffman tree. For testing purposes, we can also
 * create a Huffman code with a given HuffmanNode as the root.
 * 
 * You may use java.util.PriorityQueue rather than MyPriorityQueue.
 * On the other hand, exercising your PQ implementation is one way
 * to make sure that it is correct! 
 * 
 * All of the constructors should create the HuffmanTree and
 * the data structures necessary for efficient encoding and decoding.
 * This means that the constructor that takes frequency data should
 * build information for decoding even though that would never
 * happen in an actual implementation (you decode by first reading
 * in the header, not counting the frequencies). Doing this allows
 * you and the staff better ways to write unit tests.
 */
public class HuffmanCode {

	/** Code bit for a leaf node in file-based tree representation */
	private static final int LEAF = 0;
	/** Code bit for a parent node in file-based tree representation */
	private static final int PARENT = 1;

	/** Code bit for the left child in the file-based tree representation */
	private static final int LEFT = 0;
	/** Code bit for the right child in the file-based tree representation */
	private static final int RIGHT = 1;
	
	public HuffmanNode root;
	public MyPriorityQueue<HuffmanNode> huffQueue;
	
	/**
	 * Creates a HuffmanCode given a Huffman tree.
	 * 
	 * @throws NullPointerException
	 *             if root is null
	 */
	public HuffmanCode(HuffmanNode root) {
		if(root == null)
			throw new NullPointerException();
		this.root = root;
	}

	/**
	 * <p>
	 * Reads the Huffman header in from br and deserializes the data at the leafs of
	 * the tree with br.readByte(). The data format for this header is defined
	 * by <tt>writeHeader</tt>
	 * </p>
	 * 
	 * <p>
	 * Note that this is not used to read in a file you want to compress, but is
	 * used to read-in the Huffman codes from the header of an already
	 * compressed file.
	 * </p>
	 * 
	 * @throws IOException
	 *             If there is a problem reading from the bit reader, if the
	 *             file ends before the full header can be read, or if the
	 *             header is not valid.
	 */
	public HuffmanCode(BitReader br) throws IOException 
	{
		this.root = bitReaderToTree(br);
	}
	
	/*
	 * recursively read the bitreader, and convert it into a tree
	 * returns the root of the tree
	 */
	public HuffmanNode bitReaderToTree(BitReader br) throws IOException
	{
		int bit = br.readBit();

		if(bit == PARENT)
		{
			HuffmanNode left = bitReaderToTree(br);
			HuffmanNode right = bitReaderToTree(br);
			return new HuffmanNode(left,right);
		}
		else if (bit == LEAF)
		{	
			return new HuffmanNode((byte)br.readByte());
		}
		else //if bit is -1 or something happened to stream..
			throw new IOException();
	}

	/**
	 * Takes a list of (Byte, Frequency) pairs (here represented as a map) and
	 * builds a tree for encoding the data items using the Huffman algorithm.
	 * 
	 * @throws NullPointerException
	 *             If freqs is null
	 * @throws IllegalArgumentException
	 *             if freqs is empty
	 */
	public HuffmanCode(Map<Byte, Integer> freqs) 
	{
		if(freqs == null)
			throw new NullPointerException();
		if(freqs.size() == 0)
			throw new IllegalArgumentException();
		
		Comparator<HuffmanNode> hNodeComparator = new Comparator<HuffmanNode>()
		{
			@Override
			public int compare(HuffmanNode o1, HuffmanNode o2) {
				return o1.getFreq() - o2.getFreq();
			}	
		};
		
		huffQueue = new MyPriorityQueue<HuffmanNode>(hNodeComparator);
		
		// iterate through the map, and populate a priority queue out of huffman nodes
		for(Map.Entry<Byte, Integer> entry : freqs.entrySet())
		{
			huffQueue.offer(new HuffmanNode(entry.getValue(),entry.getKey()));
		}
		
		/*
		 * compress the queue:
		 * pop the 2 smallest elements, build them into a new node, and offer the new node back into the queue
		 */
		while(huffQueue.size() > 1)
		{
			huffQueue.offer(new HuffmanNode(huffQueue.poll(), huffQueue.poll()));
		}
		
		// update the root after compressing
		root = huffQueue.peek();
	}

	/**
	 * <p>
	 * Turns this Huffman code into a stream of bits suitable for including in a
	 * compressed file.
	 * </p>
	 * 
	 * <p>
	 * The format for the tree is defined recursively. To emit the entire tree,
	 * you start by emitting the root. When you emit a node, if the node is a
	 * leaf node, you write the bit <tt>LEAF</tt> and then call the
	 * <tt>writeByte</tt> method of <tt>BitWriter</tt> on the nodes value.
	 * Otherwise, you emit the bit <tt>PARENT</tt>, then emit the left and
	 * right node.
	 * </p>
	 * 
	 * @param writer
	 *            A bit writer to write to
	 * @throws NullPointerException
	 *             if w is null
	 * @throws IOException
	 *             If there is a problem writing to the underlying stream
	 */
	public void writeHeader(BitWriter writer) throws IOException {
		if(writer == null)
			throw new NullPointerException();
		
		printHeader(this.root,writer);
	}

	/*
	 * recursively write the header file from the root
	 */
	public void printHeader(HuffmanNode root, BitWriter writer) throws IOException
	{
		// first write the root
		if(root.getLeft() != null && root.getRight() != null)
		{
			writer.writeBit(PARENT);
		}
		
		//traverse left edge..
		if(root.getLeft() != null)
		{
			printHeader(root.getLeft(),writer);
		}
		
		// write leaf if encountered
		if(root.getLeft() == null && root.getRight() == null)
		{
			writer.writeBit(LEAF);
			writer.writeByte(root.getValue());
		}
		
		//traverse right edge
		if(root.getRight() != null)
		{
			printHeader(root.getRight(),writer);
		}
	}
	
	/**
	 * This method reads bits from the reader until the next codeword (from the
	 * given Reader) has been read in. It returns the byte that the code
	 * corresponds to. The data format for this is defined by <tt>encode</tt>
	 * 
	 * @param r
	 *            BitReader to read in the next codeword from
	 * @throws IOException
	 *             If there is an I/O error in the underlying reader. Also, if
	 *             the file contains invalid data or ends unexpectedly
	 * @throws NullPointerException
	 *             if r is null
	 * @return The data object read in
	 */
	public Byte decode(BitReader r) throws IOException {
		if(r == null)
			throw new NullPointerException();
		
		if(root == null)
			throw new NullPointerException();

		int decodeBit;

		HuffmanNode root = this.root;

		//loop through the bitreacher one bit at a time and traverse to the correct end node.
		while((decodeBit = r.readBit()) > -1)
		{
			if(decodeBit == LEFT && root.getLeft() != null)
			{
					root = root.getLeft();
			}
			else if(decodeBit == RIGHT  && root.getRight() != null)
			{	
					root = root.getRight();
			}
			if(root.getLeft() == null && root.getRight() == null && root.getValue() != 0)
			{
				return root.getValue();
			}
		}
		
		if(root.getLeft() == null && root.getRight() == null && root.getValue() != 0)
		{
			return root.getValue();
		}

		//end of stream reached, and we traversed to the right node, return the value of the node
			throw new IOException();
	}

	/**
	 * This method takes a data item and emits the corresponding codeword. The bits
	 * <tt>LEFT</tt> and <tt>RIGHT</tt> are written so that if one takes
	 * that path in the Huffman tree they will get to the leaf node representing
	 * <tt>Item</tt>.
	 * 
	 * @param item
	 *            value to encode
	 * @param writer
	 *            BitWriter to write the code word (Huffman Code for that
	 *            string)
	 * @throws NullPointerException
	 *             if the item or writer is null.
	 * @throws IllegalArgumentException
	 *             if the item doesn't exist in this huffman coding
	 */
	public void encode(Byte item, BitWriter writer) throws IOException 
	{
		//finds the node, returns the node, and climbs back the tree, inserts the encoding into a stack
		//then writes them onto the writer.. is there a better way?
		HuffmanNode hn = findNode(item, this.root);
		Stack<Integer> stack = new Stack<Integer>();
		
		if(hn == null)
			throw new IllegalArgumentException();
		
		else
		{
			while(hn.getParent() != null)
			{
				if(hn.getParent().getLeft() == hn)
					stack.push(LEFT);
				else if(hn.getParent().getRight() == hn)
					stack.push(RIGHT);
				
				hn = hn.getParent();
			}
			while(!stack.empty())
			{
				writer.writeBit(stack.pop());
			}
		}
	}
	
	public HuffmanNode findNode(Byte item, HuffmanNode root)
	{
		HuffmanNode find = null;
		
		//if the node is found.. return it.
		if(item.compareTo(root.getValue()) == 0)
			return root;
		
		//traverse left edge..
		if(root.getLeft() != null)
			find = findNode(item,root.getLeft());
		
		//traverse right edge
		if(root.getRight() != null && find == null)
			find = findNode(item,root.getRight());
		
		return find;
	}

	/**
	 * Gets the root of the Huffman tree. We use this for testing.
	 */
	public HuffmanNode getCodeTreeRoot() {
		return root;
	}
}
